package com.yeskery.transfer.shell;

import com.yeskery.transfer.config.DataBaseProperties;
import com.yeskery.transfer.core.*;
import com.yeskery.transfer.data.*;
import com.yeskery.transfer.util.TimeUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.shell.Availability;
import org.springframework.shell.standard.ShellComponent;
import org.springframework.shell.standard.ShellMethod;
import org.springframework.shell.standard.ShellMethodAvailability;
import org.springframework.shell.standard.ShellOption;
import org.springframework.util.StringUtils;

import javax.validation.constraints.Positive;
import java.sql.Connection;
import java.sql.SQLException;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * 转移数据的shell
 * @author shuangjiang.du
 * @date 2018-12-14 14:10
 * @version 1.0
 */
@Slf4j
@ShellComponent
public class DataTransferShell {
	/** 默认的最大线程数 */
	private static final int MAX_THREAD_COUNT = 8;

	/** 等待时间 */
	private static final int WAIT_TIME = 2000;

	/** 连接工具 */
	private final ConnectionKit connectionKit;

	/** 数据工具 */
	private final DataKit dataKit;

	/** 数据源工具 */
	private final DataSourceKit dataSourceKit;

	/** 系统资源对象 */
	private final Resource resource;

	/**
	 * 构建一个转移数据的shell
	 * @param connectionKit 连接工具
	 * @param dataKit 数据工具
	 * @param dataSourceKit 数据源工具
	 * @param resource 系统资源对象
	 */
	@Autowired
	public DataTransferShell(ConnectionKit connectionKit, DataKit dataKit, DataSourceKit dataSourceKit, Resource resource) {
		this.connectionKit = connectionKit;
		this.dataKit = dataKit;
		this.dataSourceKit = dataSourceKit;
		this.resource = resource;
	}

	/**
	 * 从服务器的指定表数据传输到客户端，为保证数据转移的有效性，如果在转移过程中，因为数据原因发生异常，该执行线程会中止，
	 * 但其他线程不会停止，可以修正错误数据后，再重新转移数据不完整的表数据。
	 * @param tablePattern 表的正则表达式，默认为所有表。
	 * @param excludePattern 从已经匹配到的表列表中移除指定表。
	 * @param force <code>Boolean</code>值，默认为<code>false</code>。是否强制传输，如果该值为<code>true</code>，
	 *              如果客户端也存在这些表，将会先删除客户端的这些表，再重新创建表进行数据传输。
	 * @param single <code>Boolean</code>值，默认为<code>false</code>。是否以单表单线程方式运行，即使在单表数据量
	 *               较大的情况下也使用单线程，如果该值为<code>false</code>，参数<code>pagination</code>将会自动失效，
	 *               如果开启多线程，将会按照参数<code>pagination</code>配置的数量开启多线程方式传输单表数据。
	 * @param pagination 该值针对单表数据量较大的情况，如果数据量较大将会根据当前CPU数量和总页数开启多线程传输，
	 *                   以提高传输速率，如果不想开启该功能，可以将参数<code>single</code>设置为<code>true</code>，
	 *                   该参数的默认值为{@link Page#DEFAULT_PAGINATION_NUM}。
	 * @see Page#getThreadCount()
	 */
	@ShellMethod(value = "Transfer data from the server to the client for the specified table."
			, key = {"data-tran", "data-transfer"})
	public void transferData(@ShellOption(value = {"--table", "--t"}, defaultValue = ".*") String tablePattern,
			@ShellOption(value = {"--exclude", "--e"}, defaultValue = "") String excludePattern,
			@ShellOption(value = {"--force", "--f"}, defaultValue = "false") String force,
			@ShellOption(value = {"--single", "--s"}, defaultValue = "false") String single,
			@ShellOption(value = {"--pagination", "--p"}, defaultValue = Page.DEFAULT_PAGINATION_NUM + "") @Positive Integer pagination) {
		DataBaseProperties dataBaseProperties = resource.get(KeyList.DATA_CONFIG_KEY, DataBaseProperties.class);
		ThreadPoolTaskExecutor taskExecutor = null;
		long startTime = System.currentTimeMillis();
		try (Connection serverConnection = connectionKit.getConnection(dataBaseProperties.getServerTransferProperties());
			 Connection clientConnection = connectionKit.getConnection(dataBaseProperties.getClientTransferProperties())){
			List<String> tableNameList = getPatternTableList(serverConnection, tablePattern);
			if (tableNameList == null || tableNameList.size() < 1) {
				System.out.println("Not matched table[" + tablePattern + "]");
				System.out.println(new Result().createBlankLine());
				return;
			}
			if (StringUtils.hasText(excludePattern)) {
				Pattern pattern = Pattern.compile(excludePattern);
				tableNameList = tableNameList.stream().filter(
						tableName -> !pattern.matcher(tableName).matches()).collect(Collectors.toList());
			}
			if (tableNameList.size() < 1) {
				System.out.println("Not matched table[" + tablePattern + "]");
				System.out.println(new Result().createBlankLine());
				return;
			}
			taskExecutor = doTransferData(dataBaseProperties, serverConnection, clientConnection, tableNameList
					, Boolean.parseBoolean(force), Boolean.parseBoolean(single), pagination, startTime);
		} catch (SQLException | ClassNotFoundException e) {
			throw new SproutException("Data table does not exist or database connection fails.", e);
		} catch (InterruptedException e) {
			throw new SproutException("Blocking the current thread exception.", e);
		} finally {
			if (taskExecutor != null) {
				taskExecutor.shutdown();
			}
		}
	}

	/**
	 * 检查客户端和管理端的配置信息有效性
	 * @return Availability对象
	 */
	@ShellMethodAvailability({"data-tran", "data-transfer"})
	public Availability databaseCommandAvailability() {
		return resource.get(KeyList.CLIENT_CONFIG_VALID, false)
				&& resource.get(KeyList.SERVER_CONFIG_VALID, false) ? Availability.available() :
				Availability.unavailable("Set up the database configuration information first.");
	}

	/**
	 * 执行数据转移的方法
	 * @param dataBaseProperties 数据库配置对象
	 * @param serverConnection 服务器连接对象
	 * @param clientConnection 客户端连接对象
	 * @param tableNameList 表名集合
	 * @param force 是否强制转移
	 * @param single 是否以单表单线程方式启动转移
	 * @param pagination 如果以多表进行转移，每条线程需要执行的页码数
	 * @param startTime 执行该功能开始的时间戳
	 * @return 线程池对象
	 * @throws SQLException SQLException
	 * @throws InterruptedException InterruptedException
	 */
	private ThreadPoolTaskExecutor doTransferData(DataBaseProperties dataBaseProperties, Connection serverConnection
			, Connection clientConnection, List<String> tableNameList, boolean force, boolean single
			, Integer pagination, long startTime)
			throws SQLException, InterruptedException {
		if (force) {
			createForceTables(serverConnection, clientConnection, tableNameList);
			System.out.println(new Result().createBlankLine());
		}
		//初始化连接池
		SproutDataSource serverDataSource = new SproutDataSource(dataBaseProperties.getServerTransferProperties()
				, connectionKit);
		SproutDataSource clientDataSource = new SproutDataSource(dataBaseProperties.getClientTransferProperties()
				, connectionKit);
		resource.add(KeyList.SERVER_DATA_SOURCE, serverDataSource);
		resource.add(KeyList.CLIENT_DATA_SOURCE, clientDataSource);
		//计算需要的线程数
		Map<String, Page> pageMap = new HashMap<>(tableNameList.size());
		Map<String, Counter> counterMap = new HashMap<>(tableNameList.size());
		int threadCount, transferCount = 0;
		for (String tableName : tableNameList) {
			Page page = dataKit.getPage(serverDataSource.getConnection(), tableName);
			if (pagination != Page.DEFAULT_PAGINATION_NUM && pagination > 0) {
				page.setPaginationNum(pagination);
			}
			pageMap.put(tableName, page);
			transferCount += page.getThreadCount();
			counterMap.put(tableName, new Counter());
		}
		if (!single) {
			if (transferCount > MAX_THREAD_COUNT) {
				threadCount = MAX_THREAD_COUNT;
			} else {
				threadCount = transferCount;
			}
		} else {
			if (tableNameList.size() > MAX_THREAD_COUNT) {
				threadCount = MAX_THREAD_COUNT;
			} else {
				threadCount = tableNameList.size();
			}
		}
		//初始化线程池对象
		ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
		taskExecutor.setCorePoolSize(threadCount);
		taskExecutor.setMaxPoolSize((int) (threadCount * 1.5));
		taskExecutor.initialize();
		CountDownLatch countDownLatch =  new CountDownLatch(transferCount);
		Result result = new Result();
		NumberFormat numberFormat = NumberFormat.getPercentInstance();
		//为每个表开启线程进行数据传输
		startTableThread(tableNameList, taskExecutor, pageMap, counterMap, result, serverDataSource
				, clientDataSource, countDownLatch, numberFormat, single);
		//主线程进行阻塞等待各子线程执行完成
		countDownLatch.await();
		//打印输出报告
		printTransferReport(result, tableNameList, pageMap, counterMap, numberFormat, startTime);
		return taskExecutor;
	}
	/**
	 * 从指定的数据库连接获取符合指定的正则表达式的表列表
	 * @param connection 数据库连接
	 * @param tablePattern 表的正则表达式
	 * @return 匹配正则表达式的表列表
	 * @throws SQLException SQLException
	 */
	private List<String> getPatternTableList(Connection connection, String tablePattern) throws SQLException {
		List<String> tableNameList = dataKit.getTableNames(connection);
		Pattern pattern = Pattern.compile(tablePattern);
		return  tableNameList.stream().filter(
				tableName -> pattern.matcher(tableName).matches()).collect(Collectors.toList());
	}

	/**
	 * 强制创建表的方法
	 * @param serverConnection 服务器端连接
	 * @param clientConnection 客户端连接
	 * @param tableNameList 表明集合
	 * @throws SQLException SQLException
	 */
	private void createForceTables(Connection serverConnection, Connection clientConnection, List<String> tableNameList)
			throws SQLException {
		for (String tableName : tableNameList) {
			System.out.println("Start creating table[" + tableName + "]...");
			List<MetaData> tableMetaDataList = dataKit.getMetaData(serverConnection, tableName);
			dataKit.deleteTable(clientConnection, tableName);
			dataKit.createTable(clientConnection, tableName, tableMetaDataList);
			System.out.println("Table[" + tableName + " Create success");
		}
	}

	/**
	 * 开始转移数据的方法
	 * @param clientDataSource 客户端数据源
	 * @param tableName 表名
	 * @param dataLists 数据集合
	 * @param counter 计数器
	 * @param numberFormat 数字格式化对象
	 * @param page 页码对象
	 */
	private void startTransferData(SproutDataSource clientDataSource, String tableName, List<List<Data>> dataLists,
			Counter counter, NumberFormat numberFormat, Page page) {
		dataLists.forEach(dataList -> {
			try {
				dataKit.insertData(dataSourceKit.getConnection(clientDataSource), tableName, dataList);
				long count = counter.increment();
				if ((System.currentTimeMillis() - counter.getTimestamp()) >= WAIT_TIME) {
					counter.update();
					System.out.println("Table[" + tableName + "] transmission completed "
							+ numberFormat.format(((double) count) / page.getCount()));
				}
			} catch (SQLException e) {
				log.error("Error execution table: {}, data: {}", tableName, dataList);
				log.error("Error Stack Information", e);
				throw new SproutException("Data table does not exist or database connection fails.", e);
			} catch (InterruptedException e) {
				throw new SproutException("Blocking the current thread exception.", e);
			}
		});
	}

	/**
	 * 为每个表开启线程进行数据传输
	 * @param tableNameList 需要进行数据传输的表名集合
	 * @param taskExecutor 线程池
	 * @param pageMap 表页码Map
	 * @param counterMap 表计数器Map
	 * @param result 输出结果对象
	 * @param serverDataSource 服务器数据源
	 * @param clientDataSource 客户端数据源
	 * @param countDownLatch 线程计数器
	 * @param numberFormat 数字格式化器
	 * @param single 是否以单线程传输单表数据
	 */
	private void startTableThread(List<String> tableNameList, ThreadPoolTaskExecutor taskExecutor
			, Map<String, Page> pageMap, Map<String, Counter> counterMap, Result result
			, SproutDataSource serverDataSource, SproutDataSource clientDataSource
			, CountDownLatch countDownLatch, NumberFormat numberFormat, boolean single) {
		for (String tableName : tableNameList) {
			taskExecutor.execute(() -> {
				try {
					System.out.println(result.createBlankLine());
					System.out.println(result.createContentLine("Start transmission from[" + tableName + "]..."));
					List<MetaData> metaDataList = dataKit.getMetaData(dataSourceKit.getConnection(serverDataSource), tableName);
					List<String> filedList = metaDataList.stream().map(MetaData::getName)
							.collect(Collectors.toList());
					Page page = pageMap.get(tableName);
					Counter counter = counterMap.get(tableName);
					if (single) {
						for (int i = 1; i <= page.getPages(); i++) {
							List<List<Data>> dataLists = getPageDataList(serverDataSource, filedList, tableName, i);
							startTransferData(clientDataSource, tableName, dataLists, counter, numberFormat, page);
						}
					} else {
						//以多线程方式启动数据传输
						multithreadingTransfer(page, taskExecutor, serverDataSource, clientDataSource,tableName
								, filedList, counter, numberFormat, countDownLatch);
					}
					System.out.println(result.createBlankLine());
					System.out.println(result.createContentLine("Table[" + tableName + "] successful transmission"));
				} catch (SQLException e) {
					throw new SproutException("Data table does not exist or database connection fails.", e);
				} catch (InterruptedException e) {
					throw new SproutException("Blocking the current thread exception.", e);
				} finally {
					if (single) {
						countDownLatch.countDown();
					}
				}
			});
		}
	}

	/**
	 * 以多线程方式启动数据传输
	 * @param page 页码对象
	 * @param taskExecutor 线程池对象
	 * @param serverDataSource 服务器端数据源
	 * @param clientDataSource 客户端数据源
	 * @param tableName 表名
	 * @param filedList 字段集合
	 * @param counter 计数器
	 * @param numberFormat 数字格式化器
	 * @param countDownLatch 线程计数器
	 */
	private void multithreadingTransfer(Page page, ThreadPoolTaskExecutor taskExecutor, SproutDataSource serverDataSource
			,SproutDataSource clientDataSource, String tableName, List<String> filedList, Counter counter
			, NumberFormat numberFormat, CountDownLatch countDownLatch) {
		long[] pages = page.getThreadPage(page.getThreadCount());
		long[] startPages = page.getStartPage(pages);
		for (int i = 0;i < startPages.length;i++) {
			long endPage, startPage = startPages[i];
			if (i == startPages.length - 1) {
				endPage = page.getPages();
			} else {
				endPage = startPages[i + 1] - 1;
			}
			taskExecutor.execute(() -> {
				try {
					for (int j = (int) startPage; j <= endPage; j++) {
						List<List<Data>> dataLists = getPageDataList(serverDataSource, filedList, tableName, j);
						startTransferData(clientDataSource, tableName, dataLists, counter, numberFormat, page);
					}
				} catch (SQLException e) {
					throw new SproutException("Data table does not exist or database connection fails.", e);
				} catch (InterruptedException e) {
					throw new SproutException("Blocking the current thread exception.", e);
				} finally {
					countDownLatch.countDown();
				}
			});
		}
	}

	/**
	 * 获取指定页的数据集合的方法
	 * @param serverDataSource 服务器数据源
	 * @param filedList 字段集合
	 * @param tableName 表名
	 * @param page 页码
	 * @return 指定页的数据集合的方法
	 * @throws SQLException SQLException
	 * @throws InterruptedException InterruptedException
	 */
	private List<List<Data>> getPageDataList(SproutDataSource serverDataSource, List<String> filedList
			, String tableName, int page) throws SQLException, InterruptedException {
		return dataKit.getData(dataSourceKit.getConnection(serverDataSource), resultSet -> {
			List<List<Data>> dataValueList = new ArrayList<>();
			while (resultSet.next()) {
				List<Data> dataList = new ArrayList<>();
				filedList.forEach(filed -> {
					Data data = new Data();
					data.setField(filed);
					try {
						data.setData(resultSet.getObject(filed));
					} catch (SQLException e) {
						throw new SproutException("Data table does not exist or database connection fails.", e);
					}
					dataList.add(data);
				});
				dataValueList.add(dataList);
			}
			return dataValueList;
		}, tableName, page);
	}

	/**
	 * 打印转移报告
	 * @param result 输出对象
	 * @param tableNameList 表名集合
	 * @param pageMap 表页码Map
	 * @param counterMap 表计数器Map
	 * @param numberFormat 数字格式化器
	 */
	private void printTransferReport(Result result, List<String> tableNameList, Map<String, Page> pageMap
			, Map<String, Counter> counterMap, NumberFormat numberFormat, long startTime) {
		System.out.println(result.createContentLine("All table transfers completed"));
		System.out.println(result.createContentLine("Data Transfer Report"));
		System.out.println("Transmission time: "
				+ TimeUtils.toUnitString((System.currentTimeMillis() - startTime) / 1000));
		System.out.println("Number of tables transmitted: " + tableNameList.size());
		long totalRowCount = pageMap.values().stream().mapToLong(Page::getCount).sum();
		System.out.println("Total number of records transmitted: " + totalRowCount);
		long totalSuccessRowCount = counterMap.values().stream().mapToLong(Counter::get).sum();
		System.out.println("Number of successful records: " + totalSuccessRowCount);
		System.out.println("Success rate: " + numberFormat.format(totalSuccessRowCount / totalRowCount));
		System.out.println(result.createBlankLine());
	}
}
